package middleware.optimize;

import middleware.llvm.Module;
import middleware.llvm.instruction.*;
import middleware.llvm.value.FunctionIRType;

import java.util.HashSet;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * 死代码消除
 * @author 杨凯伟
 */
public class DeadCodeElimination {
    private final Module irModule;

    public DeadCodeElimination(Module irModule) {
        this.irModule = irModule;
    }

    public Module pass() {
        irModule.getFunctions().forEach(this::passFunc);
        return irModule;
    }

    private void passFunc(FunctionIRType func) {
//        Set<Instruction> initialUsefulSet = new HashSet<>();

//        for (var block : func.getBasicBlocks()) {
//            for (var inst : block.getInstructions()) {
//                if (inst instanceof ReturnInstruction
//                        || inst instanceof BrInstruction
//                        || inst instanceof CallInstruction
//                        || inst instanceof StoreInstruction) {
//                    initialUsefulSet.add(inst);
//                }
//            }
//        }
        Set<Instruction> initialUsefulSet = func.getBasicBlocks().stream()
                .flatMap(block -> block.getInstructions().stream())
                .filter(inst -> inst instanceof ReturnInstruction
                        || inst instanceof BrInstruction
                        || inst instanceof CallInstruction
                        || inst instanceof StoreInstruction)
                .collect(Collectors.toSet());

        var usefulSet = getUsefulClosure(initialUsefulSet);
        for (var block : func.getBasicBlocks()) {
            var instructions = block.getInstructions();
            for (int i = 0; i < instructions.size(); i++) {
                var inst = instructions.get(i);
                if (!usefulSet.contains(inst) && !(inst instanceof AllocaInstruction)) {
                    inst.replaceAllUseWith(null, false);
                    i--;
                }
            }
        }

    }

    private Set<Instruction> getUsefulClosure(Set<Instruction> instSet) {
        Set<Instruction> usefulSet = new HashSet<>();
        instSet.forEach(inst -> usefulSet.addAll(getUsefulClosure(inst)));
        return usefulSet;
    }

    private Set<Instruction> getUsefulClosure(Instruction inst) {
        Set<Instruction> usefulSet = new HashSet<>();
        usefulSet.add(inst);
//        for (var operand : inst.getOperands()) {
//            if (operand instanceof Instruction instOperand) {
//                usefulSet.addAll(getUsefulClosure(instOperand));
//            }
//        }
        inst.getOperands().stream()
                .filter(operand -> operand instanceof Instruction)
                .map(operand -> (Instruction) operand)
                .forEach(instOperand -> usefulSet.addAll(getUsefulClosure(instOperand)));
        return usefulSet;
    }
}
